require "test/unit/ui/console/testrunner"

module Test
  module Unit
    module UI
      module Console
        class ColorizedTestRunner < TestRunner
          extend TestRunnerUtilities

          SCHEMES = {
            :default => {
              "success" => Color.new("green", :bold => true),
              "failure" => Color.new("red", :bold => true),
              "pending" => Color.new("magenta", :bold => true),
              "omission" => Color.new("blue", :bold => true),
              "notification" => Color.new("cyan", :bold => true),
              "error" => Color.new("yellow", :bold => true),
            },
          }

          def initialize(suite, output_level=NORMAL, io=STDOUT)
            super
            @use_color = guess_color_availability
            @color_scheme = SCHEMES[:default]
            @reset_color = Color.new("reset")
          end

          private
          def add_fault(fault)
            @faults << fault
            output_single_with_color(fault.single_character_display,
                                     fault_color(fault),
                                     PROGRESS_ONLY)
            @already_outputted = true
          end

          def test_finished(name)
            unless @already_outputted
              output_single_with_color(".",
                                       @color_scheme["success"],
                                       PROGRESS_ONLY)
            end
            nl(VERBOSE)
            @already_outputted = false
          end

          def finished(elapsed_time)
            nl
            output("Finished in #{elapsed_time} seconds.")
            @faults.each_with_index do |fault, index|
              nl
              output_single("%3d) " % (index + 1))
              output_with_color(fault.long_display, fault_color(fault))
            end
            nl
            output_with_color(@result.to_s, result_color)
          end

          def fault_color(fault)
            @color_scheme[fault.class.name.split(/::/).last.downcase]
          end

          def result_color
            if @result.passed?
              if @result.pending_count > 0
                @color_scheme["pending"]
              elsif @result.notification_count > 0
                @color_scheme["notification"]
              else
                @color_scheme["success"]
              end
            elsif @result.error_count > 0
              @color_scheme["error"]
            elsif @result.failure_count > 0
              @color_scheme["failure"]
            end
          end

          def output_with_color(message, color=nil, level=NORMAL)
            return unless output?(level)
            output_single_with_color(message, color, level)
            @io.puts
          end

          def output_single_with_color(message, color=nil, level=NORMAL)
            return unless output?(level)
            if @use_color and color
              message = "%s%s%s" % [color.escape_sequence,
                                    message,
                                    @reset_color.escape_sequence]
            end
            @io.write(message)
            @io.flush
          end

          def guess_color_availability
            return false unless @io.tty?
            term = ENV["TERM"]
            return true if term and (/term\z/ =~ term or term == "screen")
            return true if ENV["EMACS"] == "t"
            false
          end
        end
      end
    end

    AutoRunner::RUNNERS[:console] = Proc.new do
      Test::Unit::UI::Console::ColorizedTestRunner
    end
  end
end
